/**
 * @license
 * Copyright 2018 Google Inc. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import { sequences } from '@magenta/music'
const { quantizeNoteSequence, clone } = sequences

const STEPS_PER_BAR = 16
const STEPS_PER_QUARTER = 4

/**
 * Reconstructs unquantized sequence by breaking it into as large of chunks as
 * possible that are supported by the given models.
 *
 * @param inSeq the NoteSequence to reconstruct.
 * @param models a list of MusicVAE models where the value of `index + 1`
 * specifies the number of bars that the model supports for reconstruction.
 * @param temperature the sampling temperature to use in reconstruction.
 **/
export async function reconstructBySize(inSeq, models, temperature=1){
	inSeq = quantizeNoteSequence(inSeq, STEPS_PER_QUARTER)

	// Process in as large of chunks as possible.
	const maxChunkSize = models.length * STEPS_PER_BAR
	const secondsPerStep = 1 / sequences.stepsPerQuarterToStepsPerSecond(
		STEPS_PER_QUARTER, inSeq.tempos[0].qpm)
	const outputs = []
	for (let startOffset = 0; startOffset < inSeq.totalQuantizedSteps; startOffset+=maxChunkSize){
		const chunk = clone(inSeq)
		const endOffset = Math.min(startOffset + maxChunkSize, inSeq.totalQuantizedSteps)
		chunk.notes = inSeq.notes
			.map(n => Object.assign({}, n))
			.filter(n => startOffset <= n.quantizedStartStep && n.quantizedStartStep < endOffset)
			.map(n => {
				n.startTime -= startOffset * secondsPerStep
				n.endTime -= startOffset * secondsPerStep
				n.quantizedStartStep -= startOffset
				n.quantizedEndStep -= startOffset
				return n
			})
		chunk.totalQuantizedSteps = endOffset - startOffset
		chunk.totalTime = chunk.totalQuantizedSteps * secondsPerStep

		// Select model based on the number of actual bars in the chunk.
		const numBars = Math.ceil(chunk.totalQuantizedSteps / STEPS_PER_BAR)
		const modelIndex = numBars - 1
		const z = await models[modelIndex].encode([chunk])
		const output = await models[modelIndex].decode(z, temperature, undefined, undefined, inSeq.tempos[0].qpm)
		z.dispose()
		outputs.push(output[0])
	}
	const reconstruction = concat(...outputs)
	return reconstruction
}

function concat(...args){
	if (args.length === 2){
		const [seqA, seqB] = args
		const outputSequence = clone(seqA)
		seqB.notes.forEach(note => {
			const clonedNote = Object.assign({}, note)
			clonedNote.startTime += seqA.totalTime
			clonedNote.endTime += seqA.totalTime
			outputSequence.notes.push(clonedNote)
		})
		outputSequence.totalTime = seqA.totalTime + seqB.totalTime
		return outputSequence
	} else if (args.length > 2){
		// Recurse.
		const first = args.shift()
		return concat(first, concat(...args))
	} else {
		// Nothing to concat.
		return args[0]
	}
}
